"""
TensorFlow MPS 诊断和性能测试工具

这个脚本检查TensorFlow在Mac上的MPS支持情况，并进行性能测试
"""

import tensorflow as tf
import time
import platform
import sys
import numpy as np
import matplotlib.pyplot as plt

def diagnose_environment():
    """
    诊断环境配置
    """
    print("=== 环境诊断 ===")
    print(f"Python版本: {sys.version}")
    print(f"操作系统: {platform.platform()}")
    print(f"TensorFlow版本: {tf.__version__}")
    print(f"TensorFlow编译时CUDA支持: {tf.test.is_built_with_cuda()}")
    
    # 检查是否安装了tensorflow-metal
    try:
        import tensorflow_metal
        print(f"TensorFlow Metal版本: {tensorflow_metal.__version__}")
    except ImportError:
        print("未安装tensorflow-metal插件")
    
    print()

def check_physical_devices():
    """
    检查所有物理设备
    """
    print("=== 物理设备检查 ===")
    physical_devices = tf.config.list_physical_devices()
    print("所有物理设备:")
    for device in physical_devices:
        print(f"  {device}")
    
    print(f"\n设备详细信息:")
    for device in physical_devices:
        try:
            details = tf.config.experimental.get_device_details(device)
            print(f"  {device}: {details}")
        except Exception as e:
            print(f"  {device}: 无法获取详细信息 - {e}")
    
    gpu_devices = tf.config.list_physical_devices('GPU')
    print(f"\nGPU设备数量: {len(gpu_devices)}")
    
    if gpu_devices:
        print("GPU设备列表:")
        for gpu in gpu_devices:
            print(f"  {gpu}")
    else:
        print("未检测到GPU设备")
    
    print()
    return len(gpu_devices) > 0

def test_mps_availability():
    """
    测试MPS可用性
    """
    print("=== MPS可用性测试 ===")
    
    # 检查MPS是否可用（TensorFlow 2.5+）
    try:
        mps_available = tf.config.list_physical_devices('GPU') and \
                       any('mps' in str(device).lower() for device in tf.config.list_physical_devices('GPU'))
        if mps_available:
            print("检测到MPS设备")
            return True
    except Exception as e:
        print(f"MPS检测异常: {e}")
    
    # 更通用的检查方式
    gpu_devices = tf.config.list_physical_devices('GPU')
    if gpu_devices:
        print("检测到GPU设备，可能支持MPS")
        return True
    else:
        print("未检测到GPU设备，MPS可能不可用")
        return False

def configure_gpu():
    """
    配置GPU设置
    """
    print("=== GPU配置 ===")
    gpu_devices = tf.config.list_physical_devices('GPU')
    if gpu_devices:
        try:
            # 启用内存增长
            for gpu in gpu_devices:
                tf.config.experimental.set_memory_growth(gpu, True)
                print(f"为 {gpu} 启用内存增长")
            
            # 设置虚拟设备（用于测试）
            # tf.config.set_logical_device_configuration(
            #     gpu_devices[0],
            #     [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
            # logical_gpus = tf.config.list_logical_devices('GPU')
            # print(f"逻辑GPU设备数量: {len(logical_gpus)}")
            
        except RuntimeError as e:
            print(f"GPU配置错误: {e}")
    else:
        print("无GPU设备需要配置")
    
    print()

def matrix_computation_test():
    """
    矩阵计算性能测试
    """
    print("=== 矩阵计算性能测试 ===")
    
    # 测试矩阵大小
    sizes = [1000, 2000, 3000]
    
    gpu_available = len(tf.config.list_physical_devices('GPU')) > 0
    
    for size in sizes:
        print(f"\n测试 {size}x{size} 矩阵乘法:")
        
        # CPU测试
        with tf.device('/CPU:0'):
            # 创建矩阵
            a = tf.random.normal([size, size])
            b = tf.random.normal([size, size])
            
            # 预热
            _ = tf.matmul(a, b)
            
            # 实际测试
            start_time = time.time()
            c = tf.matmul(a, b)
            _ = tf.reduce_sum(c)  # 确保计算完成
            cpu_time = time.time() - start_time
            
        print(f"  CPU耗时: {cpu_time:.4f} 秒")
        
        # GPU测试
        if gpu_available:
            with tf.device('/GPU:0'):
                # 创建矩阵
                a = tf.random.normal([size, size])
                b = tf.random.normal([size, size])
                
                # 预热
                _ = tf.matmul(a, b)
                
                # 实际测试
                start_time = time.time()
                c = tf.matmul(a, b)
                _ = tf.reduce_sum(c)  # 确保计算完成
                gpu_time = time.time() - start_time
                
            print(f"  GPU耗时: {gpu_time:.4f} 秒")
            print(f"  加速比: {cpu_time/gpu_time:.2f}x")
        else:
            print("  GPU测试跳过（无GPU设备）")
    
    print()

def simple_neural_network_test():
    """
    简单神经网络测试
    """
    print("=== 简单神经网络测试 ===")
    
    # 创建简单数据集
    x_train = tf.random.normal([10000, 784])
    y_train = tf.random.uniform([10000], maxval=10, dtype=tf.int32)
    
    def create_model():
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(512, activation='relu', input_shape=(784,)),
            tf.keras.layers.Dense(256, activation='relu'),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        return model
    
    # CPU测试
    print("\n在CPU上训练神经网络...")
    with tf.device('/CPU:0'):
        model_cpu = create_model()
        model_cpu.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
        
        start_time = time.time()
        model_cpu.fit(x_train, y_train, epochs=3, batch_size=128, verbose=0)
        cpu_time = time.time() - start_time
    
    print(f"CPU训练时间: {cpu_time:.2f} 秒")
    
    # GPU测试
    gpu_devices = tf.config.list_physical_devices('GPU')
    if gpu_devices:
        print("\n在GPU上训练神经网络...")
        with tf.device('/GPU:0'):
            model_gpu = create_model()
            model_gpu.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
            
            # 预热
            model_gpu.fit(x_train[:1000], y_train[:1000], epochs=1, batch_size=128, verbose=0)
            
            start_time = time.time()
            model_gpu.fit(x_train, y_train, epochs=3, batch_size=128, verbose=0)
            gpu_time = time.time() - start_time
        
        print(f"GPU训练时间: {gpu_time:.2f} 秒")
        print(f"训练加速比: {cpu_time/gpu_time:.2f}x")
    else:
        print("GPU测试跳过（无GPU设备）")
    
    print()

def advanced_neural_network_test():
    """
    高级神经网络测试
    """
    print("=== 高级神经网络测试 ===")
    
    # 创建更大的数据集
    x_train = np.random.random((100000, 512)).astype(np.float32)
    y_train = np.random.randint(0, 10, (100000,))
    
    # 创建更复杂的模型
    def create_model():
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(1024, activation='relu', input_shape=(512,)),
            tf.keras.layers.Dense(512, activation='relu'),
            tf.keras.layers.Dense(256, activation='relu'),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        return model
    
    # CPU测试
    print("\n在CPU上训练复杂神经网络...")
    with tf.device('/CPU:0'):
        model_cpu = create_model()
        model_cpu.compile(optimizer='adam', 
                         loss='sparse_categorical_crossentropy',
                         metrics=['accuracy'])
        
        start_time = time.time()
        model_cpu.fit(x_train, y_train, epochs=5, batch_size=256, verbose=1)
        cpu_training_time = time.time() - start_time
    
    print(f"CPU训练时间: {cpu_training_time:.2f} 秒")
    
    # GPU测试
    gpu_training_time = None
    if len(tf.config.list_physical_devices('GPU')) > 0:
        print("\n在GPU(MPS)上训练复杂神经网络...")
        with tf.device('/GPU:0'):
            model_gpu = create_model()
            model_gpu.compile(optimizer='adam',
                             loss='sparse_categorical_crossentropy',
                             metrics=['accuracy'])
            
            # 预热
            model_gpu.fit(x_train[:1000], y_train[:1000], epochs=1, batch_size=256, verbose=0)
            
            start_time = time.time()
            model_gpu.fit(x_train, y_train, epochs=5, batch_size=256, verbose=1)
            gpu_training_time = time.time() - start_time
        
        print(f"GPU训练时间: {gpu_training_time:.2f} 秒")
        print(f"训练加速比: {cpu_training_time/gpu_training_time:.2f}x")
    
    print()
    return cpu_training_time, gpu_training_time

def visualize_results(matrix_sizes, cpu_times, gpu_times, cpu_training_time, gpu_training_time):
    """
    可视化测试结果
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # 矩阵运算性能对比
    ax1.plot(matrix_sizes, cpu_times, 'o-', label='CPU', linewidth=2, markersize=8)
    if gpu_times:
        ax1.plot(matrix_sizes, gpu_times, 's-', label='GPU(MPS)', linewidth=2, markersize=8)
        # 计算加速比
        speedups = [c/g for c, g in zip(cpu_times, gpu_times)]
        ax2.bar(range(len(matrix_sizes)), speedups, color='green', alpha=0.7)
        ax2.set_xlabel('矩阵大小')
        ax2.set_ylabel('加速比')
        ax2.set_title('矩阵运算加速比 (CPU/GPU)')
        ax2.set_xticks(range(len(matrix_sizes)))
        ax2.set_xticklabels(matrix_sizes)
        for i, v in enumerate(speedups):
            ax2.text(i, v + 0.1, f'{v:.2f}x', ha='center', va='bottom')
    
    ax1.set_xlabel('矩阵大小')
    ax1.set_ylabel('时间 (秒)')
    ax1.set_title('矩阵运算性能对比')
    ax1.legend()
    ax1.grid(True)
    
    # 训练时间对比
    training_times = [cpu_training_time]
    labels = ['CPU']
    if gpu_training_time:
        training_times.append(gpu_training_time)
        labels.append('GPU(MPS)')
    
    ax2.bar(labels, training_times, color=['blue', 'orange'], alpha=0.7)
    ax2.set_ylabel('训练时间 (秒)')
    ax2.set_title('神经网络训练时间对比')
    for i, v in enumerate(training_times):
        ax2.text(i, v + 1, f'{v:.2f}s', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('tensorflow_mps_performance.png', dpi=300, bbox_inches='tight')
    plt.show()

def troubleshooting_guide():
    """
    故障排除指南
    """
    print("=== MPS故障排除指南 ===")
    print("如果MPS未被检测到，请尝试以下步骤:")
    print()
    print("1. 确认系统要求:")
    print("   - macOS 11.0+ (Big Sur 或更高版本)")
    print("   - 支持Metal的Mac (M1芯片或Intel芯片)")
    print()
    print("2. 安装正确的TensorFlow版本:")
    print("   pip install tensorflow-macos")
    print("   pip install tensorflow-metal")
    print()
    print("3. 检查已安装的包:")
    print("   pip list | grep tensorflow")
    print()
    print("4. 如果仍有问题，尝试:")
    print("   - 重新安装Python环境")
    print("   - 更新Xcode命令行工具")
    print("   - 检查系统是否有足够的内存")
    print()

def main():
    """
    主函数
    """
    print("TensorFlow MPS 详细诊断和性能测试")
    print("=" * 50)
    
    # 环境诊断
    diagnose_environment()
    
    # 检查物理设备
    gpu_available = check_physical_devices()
    
    # 测试MPS可用性
    mps_available = test_mps_availability()
    
    # 配置GPU
    configure_gpu()
    
    # 矩阵计算测试
    matrix_sizes = [1000, 2000, 3000]
    cpu_times = []
    gpu_times = []
    
    if gpu_available:
        # 运行矩阵测试以收集数据
        for size in matrix_sizes:
            with tf.device('/CPU:0'):
                a = tf.random.normal([size, size])
                b = tf.random.normal([size, size])
                _ = tf.matmul(a, b)
                start_time = time.time()
                c = tf.matmul(a, b)
                _ = tf.reduce_sum(c)
                cpu_time = time.time() - start_time
                cpu_times.append(cpu_time)
            
            with tf.device('/GPU:0'):
                a = tf.random.normal([size, size])
                b = tf.random.normal([size, size])
                _ = tf.matmul(a, b)
                start_time = time.time()
                c = tf.matmul(a, b)
                _ = tf.reduce_sum(c)
                gpu_time = time.time() - start_time
                gpu_times.append(gpu_time)
        
        print("=== 矩阵测试结果摘要 ===")
        for size, cpu, gpu in zip(matrix_sizes, cpu_times, gpu_times):
            print(f"{size}x{size}: CPU {cpu:.4f}s, GPU {gpu:.4f}s, 加速比 {cpu/gpu:.2f}x")
    
    # 运行神经网络测试
    print("=== 神经网络测试结果 ===")
    cpu_training_time, gpu_training_time = advanced_neural_network_test()
    
    # 显示加速比摘要
    if gpu_available and gpu_training_time:
        avg_speedup = np.mean([c/g for c, g in zip(cpu_times, gpu_times)])
        training_speedup = cpu_training_time / gpu_training_time
        print("=== 性能测试总结 ===")
        print(f"矩阵运算平均加速比: {avg_speedup:.2f}x")
        print(f"神经网络训练加速比: {training_speedup:.2f}x")
    
    # 故障排除指南
    troubleshooting_guide()
    
    print("测试完成!")

if __name__ == "__main__":
    main()
import tensorflow as tf
import time
import numpy as np
import matplotlib.pyplot as plt

def check_mps_support():
    """
    检查TensorFlow对MPS的支持
    """
    print("TensorFlow版本:", tf.__version__)
    print("是否构建支持CUDA:", tf.test.is_built_with_cuda())
    
    # 检查物理设备
    physical_devices = tf.config.list_physical_devices()
    print("\\n所有物理设备:")
    for device in physical_devices:
        print(f"  {device}")
    
    # 检查GPU设备
    gpu_devices = tf.config.list_physical_devices('GPU')
    print(f"\\nGPU设备数量: {len(gpu_devices)}")
    
    if gpu_devices:
        print("MPS支持可用")
        # 启用内存增长以避免预分配所有GPU内存
        try:
            for gpu in gpu_devices:
                tf.config.experimental.set_memory_growth(gpu, True)
        except RuntimeError as e:
            print(f"设置内存增长时出错: {e}")
        return True
    else:
        print("MPS支持不可用")
        return False

def matrix_operations_test():
    """
    矩阵运算性能测试
    """
    print("\\n矩阵运算性能测试")
    print("=" * 30)
    
    # 测试不同大小的矩阵
    matrix_sizes = [1000, 2000, 3000, 4000]
    cpu_times = []
    gpu_times = []
    
    # 检查是否有GPU可用
    gpu_available = len(tf.config.list_physical_devices('GPU')) > 0
    
    for size in matrix_sizes:
        print(f"\\n测试矩阵大小: {size}x{size}")
        
        # 在CPU上测试
        with tf.device('/CPU:0'):
            # 创建随机矩阵
            a_cpu = tf.random.normal([size, size])
            b_cpu = tf.random.normal([size, size])
            
            # 预热
            _ = tf.matmul(a_cpu, b_cpu)
            
            # 实际测试
            start_time = time.time()
            c_cpu = tf.matmul(a_cpu, b_cpu)
            _ = tf.reduce_sum(c_cpu)  # 强制计算完成
            cpu_time = time.time() - start_time
            
        cpu_times.append(cpu_time)
        print(f"  CPU时间: {cpu_time:.4f} 秒")
        
        # 在GPU上测试
        if gpu_available:
            with tf.device('/GPU:0'):
                # 创建随机矩阵
                a_gpu = tf.random.normal([size, size])
                b_gpu = tf.random.normal([size, size])
                
                # 预热
                _ = tf.matmul(a_gpu, b_gpu)
                
                # 实际测试
                start_time = time.time()
                c_gpu = tf.matmul(a_gpu, b_gpu)
                _ = tf.reduce_sum(c_gpu)  # 强制计算完成
                gpu_time = time.time() - start_time
                
            gpu_times.append(gpu_time)
            print(f"  GPU时间: {gpu_time:.4f} 秒")
            print(f"  加速比: {cpu_time/gpu_time:.2f}x")
        else:
            print("  GPU测试跳过 (MPS不可用)")
    
    return matrix_sizes, cpu_times, gpu_times

def neural_network_test():
    """
    神经网络训练性能测试
    """
    print("\\n\\n神经网络训练性能测试")
    print("=" * 30)
    
    # 创建更大的数据集
    x_train = np.random.random((100000, 512)).astype(np.float32)
    y_train = np.random.randint(0, 10, (100000,))
    
    # 创建模型
    def create_model():
        model = tf.keras.Sequential([
            tf.keras.layers.Dense(1024, activation='relu', input_shape=(512,)),
            tf.keras.layers.Dense(512, activation='relu'),
            tf.keras.layers.Dense(256, activation='relu'),
            tf.keras.layers.Dense(128, activation='relu'),
            tf.keras.layers.Dense(10, activation='softmax')
        ])
        return model
    
    # CPU测试
    print("\\n在CPU上训练神经网络...")
    with tf.device('/CPU:0'):
        model_cpu = create_model()
        model_cpu.compile(optimizer='adam', 
                         loss='sparse_categorical_crossentropy',
                         metrics=['accuracy'])
        
        start_time = time.time()
        model_cpu.fit(x_train, y_train, epochs=5, batch_size=256, verbose=1)
        cpu_training_time = time.time() - start_time
    
    print(f"CPU训练时间: {cpu_training_time:.2f} 秒")
    
    # GPU测试
    gpu_training_time = None
    if len(tf.config.list_physical_devices('GPU')) > 0:
        print("\\n在GPU(MPS)上训练神经网络...")
        with tf.device('/GPU:0'):
            model_gpu = create_model()
            model_gpu.compile(optimizer='adam',
                             loss='sparse_categorical_crossentropy',
                             metrics=['accuracy'])
            
            # 预热
            model_gpu.fit(x_train[:1000], y_train[:1000], epochs=1, batch_size=256, verbose=0)
            
            start_time = time.time()
            model_gpu.fit(x_train, y_train, epochs=5, batch_size=256, verbose=1)
            gpu_training_time = time.time() - start_time
        
        print(f"GPU训练时间: {gpu_training_time:.2f} 秒")
        print(f"训练加速比: {cpu_training_time/gpu_training_time:.2f}x")
    
    return cpu_training_time, gpu_training_time

def visualize_results(matrix_sizes, cpu_times, gpu_times, cpu_training_time, gpu_training_time):
    """
    可视化测试结果
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
    
    # 矩阵运算性能对比
    ax1.plot(matrix_sizes, cpu_times, 'o-', label='CPU', linewidth=2, markersize=8)
    if gpu_times:
        ax1.plot(matrix_sizes, gpu_times, 's-', label='GPU(MPS)', linewidth=2, markersize=8)
        # 计算加速比
        speedups = [c/g for c, g in zip(cpu_times, gpu_times)]
        ax2.bar(range(len(matrix_sizes)), speedups, color='green', alpha=0.7)
        ax2.set_xlabel('矩阵大小')
        ax2.set_ylabel('加速比')
        ax2.set_title('矩阵运算加速比 (CPU/GPU)')
        ax2.set_xticks(range(len(matrix_sizes)))
        ax2.set_xticklabels(matrix_sizes)
        for i, v in enumerate(speedups):
            ax2.text(i, v + 0.1, f'{v:.2f}x', ha='center', va='bottom')
    
    ax1.set_xlabel('矩阵大小')
    ax1.set_ylabel('时间 (秒)')
    ax1.set_title('矩阵运算性能对比')
    ax1.legend()
    ax1.grid(True)
    
    # 训练时间对比
    training_times = [cpu_training_time]
    labels = ['CPU']
    if gpu_training_time:
        training_times.append(gpu_training_time)
        labels.append('GPU(MPS)')
    
    ax2.bar(labels, training_times, color=['blue', 'orange'], alpha=0.7)
    ax2.set_ylabel('训练时间 (秒)')
    ax2.set_title('神经网络训练时间对比')
    for i, v in enumerate(training_times):
        ax2.text(i, v + 1, f'{v:.2f}s', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('/Users/jiangyong/project/heimabigdata/AI实验/tensorflow_mps_performance.png', dpi=300, bbox_inches='tight')
    plt.show()

def main():
    """
    主函数
    """
    print("TensorFlow MPS 详细性能测试")
    print("=" * 50)
    
    # 检查MPS支持
    mps_available = check_mps_support()
    
    # 矩阵运算测试
    matrix_sizes, cpu_times, gpu_times = matrix_operations_test()
    
    # 神经网络测试
    cpu_training_time, gpu_training_time = neural_network_test()
    
    # 显示总结
    print("\\n\\n测试总结")
    print("=" * 30)
    if gpu_times:
        avg_speedup = np.mean([c/g for c, g in zip(cpu_times, gpu_times)])
        print(f"矩阵运算平均加速比: {avg_speedup:.2f}x")
    
    if gpu_training_time:
        training_speedup = cpu_training_time / gpu_training_time
        print(f"神经网络训练加速比: {training_speedup:.2f}x")
    
    print("\\n测试完成!")

if __name__ == "__main__":
    main()